const WebSocket = require('ws')

class ProgressNotifier {
  constructor() {
    this.clients = new Map()
  }

  // 创建WebSocket服务器
  createServer(server) {
    this.wss = new WebSocket.Server({ server })
    
    this.wss.on('connection', (ws, req) => {
      const clientId = req.headers['sec-websocket-key']
      this.clients.set(clientId, ws)
      
      ws.on('close', () => {
        this.clients.delete(clientId)
      })
    })
  }

  // 发送进度通知
  notifyProgress(clientId, data) {
    const ws = this.clients.get(clientId)
    if (ws && ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({
        type: 'progress',
        data
      }))
    }
  }

  // 发送完成通知
  notifyComplete(clientId, data) {
    const ws = this.clients.get(clientId)
    if (ws && ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({
        type: 'complete',
        data
      }))
      ws.close()
    }
  }

  // 发送错误通知
  notifyError(clientId, error) {
    const ws = this.clients.get(clientId)
    if (ws && ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({
        type: 'error',
        error: error.message
      }))
      ws.close()
    }
  }
}

module.exports = new ProgressNotifier()